'''
下载所有股票信息，包括4列：序号、编号、是否有效和名称。
输入： 无
输出： 保存到文件 f'all_stock_info_{datetime.now():%Y%m%d}.csv'

【参考资料】
1. 技术文档地址: https://ai-cyber-security.feishu.cn/docx/AB1fdqV9xo9K9axwT3hcMOOcnzc
2. BaoStock官网: http://baostock.com/baostock/index.php/Python_API文档

'''
import os, sys
import baostock as bs
from datetime import datetime, timedelta
import pandas as pd
from utils.stock_funs import *

# 每个交易日的股票信息都可以单独下载
stock_date = f'2025-01-03'

filepath = download_all_stock_names(stock_date)
print(filepath)


# 读取数据, 默认开始日期为2010/01/01, 结束日期为当前日期
def download_stock(lg, stock_code, start_date='2000-01-01', end_date=None, forced=False):
    stock_file = f'./monthly_data/{stock_code}.csv'
    if os.path.exists(stock_file) and not forced:
        return None
    
    # 初始化参数
    if end_date == None:
        end_date= datetime.now().strftime('%Y-%m-%d')
        
    # 读取数据
    column_names = "date,open,high,low,close"
    rs = bs.query_history_k_data_plus(
        stock_code, 
        column_names, 
        start_date=start_date, 
        end_date=end_date, 
        frequency="m")

    data_list = []
    if rs.error_code == '0':
        while (rs.error_code == '0') & rs.next():
            data_list.append(rs.get_row_data())
    else:
        print('query_history_k_data_plus failed, error_code: ' + rs.error_code)
        print('query_history_k_data_plus failed, error_msg:  ' + rs.error_msg)
        return None
    
    # 保存到本地, 并返回df
    df = pd.DataFrame(data_list, columns=column_names.split(','))
    df.to_csv(stock_file, encoding="utf-8", index=False)
    return df


# 如果不存在, 则下载所有股票代码
stocknames_path = 'all_stock_names_2025-01-03.csv'
if not os.path.exists(stocknames_path):
    print('downloading...')
    stocknames_path = download_all_stock_names(f'2025-01-03')
    print(stocknames_path)


if not os.path.exists('./monthly_data'):
    os.mkdir('./monthly_data')

lg = bs.login()
if  lg.error_code != '0':
    print("Loging failed. Error code: " + lg.error_msg)
    print("Loging failed. Error message: " + lg.error_code)
    sys.exit()

# 可以指定下载的股票编码和开始日期。
#code = 'SH.603538' if len(sys.argv) < 2 else sys.argv[1].upper()
#start_date = '2010-01-01' if len(sys.argv) < 3 else sys.argv[2]



# 文件包括三列:  code, tradeStatus, code_name
df_names = pd.read_csv(stocknames_path)
start_date = '2015-01-01'
for index, row in df_names.iterrows():
    try:
        dr = row.tolist() # ['sh.000001', 1, '上证综合指数']
        download_stock(lg, dr[0], start_date)   # dr[1] =  sz.301589
        print(dr[0], dr[1], dr[2])       # 5305 sz.301589 1 诺瓦星云
    except:
        continue

